#ifndef VOXEL_H_
#define VOXEL_H_

#include <boost/json.hpp>
#include <chrono>
#include <future>
#include <mutex>
#include <random>

#include "common/basic_type.h"
#include "common/point_cloud.h"
#include "common/utility.h"
namespace autodrive {
namespace perception {

class Voxel final {
 public:
  Voxel() noexcept;
  ~Voxel();
  Voxel(Voxel const&) = delete;
  Voxel& operator=(Voxel const&) = delete;
  Voxel(Voxel&&) = delete;
  Voxel& operator=(Voxel&&) = delete;

 public:
  /// @brief voxel parameters
  struct VoxelParameters {
    /// @brief net range
    std::vector<float> net_range;
    /// @brief voxel size
    std::vector<float> voxel_size;
    /// @brief
    std::vector<int32_t> mp;
    /// @brief down stride
    int32_t sd;
    /// @brief max points per voxel
    int32_t max_points_per_voxel;
    /// @brief scale for ptq/qat
    int32_t scale;
    /// @brief multi thread
    uint16_t multi_threads;
    /// @brief is shuffle flag
    bool is_shuffle;
  };

  template <std::uint32_t maxPointSize, typename _Tp>
  bool Density(PointCloud<PointXYZIT, maxPointSize> const& pclIn,
               VoxelParameters const& parameters,
               int32_t const& maxVoxelDataSize, _Tp* voxelData,
               std::string& err) noexcept {
    auto const& voxelSize{parameters.voxel_size};
    auto const& netRange{parameters.net_range};
    auto const& m1{parameters.mp.at(0)};
    auto const& m2{parameters.mp.at(1)};
    auto const& maxPointsPerVoxel{parameters.max_points_per_voxel};

    auto const& pointCloudNum{pclIn.GetPointNum()};
    // 初始化
    auto const& ret{init(parameters, pointCloudNum, err)};
    if (!ret) {
      return false;
    }

    // voxel
    float vx{voxelSize[0] / 2 + netRange[0]};
    float vy{voxelSize[1] / 2 + netRange[1]};
    float vz{voxelSize[2] / 2 + netRange[2]};
    float z_coe{voxelSize[2] / 2};

    std::vector<int32_t> grid(voxelSize.size(), 0);

    // feature map h
    auto const& h{static_cast<int32_t>(grid_size_[1] * m2)};
    // feature map wgen_index
    auto const& w{static_cast<int32_t>(grid_size_[0] * m1)};

    // feature map stride
    auto const& stride{h * w};
    auto const& pointCloud{pclIn};
    auto doVoxelFunc{[&](int32_t startIndex, int32_t stopIndex) {
      for (auto j{startIndex}; j < stopIndex; ++j) {
        auto const& selectIndex{random_array_ptr_->at(j)};
        if (selectIndex < pointCloudNum) {
          auto& point{pointCloud.GetPoint(selectIndex)};
          grid[0] = floor((point.x - netRange[0]) / voxelSize[0]);
          if (grid[0] < 0 || grid[0] >= grid_size_[0]) {
            continue;
          }
          grid[1] = floor((point.y - netRange[1]) / voxelSize[1]);
          if (grid[1] < 0 || grid[1] >= grid_size_[1]) {
            continue;
          }
          grid[2] = floor((point.z - netRange[2]) / voxelSize[2]);
          if (grid[2] < 0 || grid[2] >= grid_size_[2]) {
            continue;
          }

          auto const& voxelPointIndex{
              gen_index(grid[0], grid[1], grid_size_[0], grid_size_[1])};

          if (voxelPointIndex >= voxel_points_num_lut_ptr_->size()) {
            continue;
          }

          auto const& voxelPointNumber{
              voxel_points_num_lut_ptr_->at(voxelPointIndex).load()};

          if (voxelPointNumber >= maxPointsPerVoxel) {
            continue;
          }

          // position in single feature map
          std::int32_t pos{(grid[1] * m1 + voxelPointNumber % m1) * w +
                           (grid[0] * m2 + voxelPointNumber / m2)};

          // feature index
          auto x_index{0 * stride + pos};
          auto y_index{1 * stride + pos};
          auto z_index{2 * stride + pos};
          auto i_index{3 * stride + pos};

          // feature value
          auto const& x{(point.x - (grid[0] * voxelSize[0] + vx)) /
                        voxelSize[0] * 2.0};
          auto const& y{(point.y - (grid[1] * voxelSize[1] + vy)) /
                        voxelSize[1] * 2.0};
          auto const& z{(point.z - vz) / z_coe};
          auto const& intensity{(static_cast<float>(point.intensity) - 127.5) /
                                127.5};

          auto const& x_b{point.x >= 0 ? 1 : -1};
          auto const& y_b{point.y >= 0 ? 1 : -1};

          if (x_index < maxVoxelDataSize) {
            voxelData[x_index] = quantify<float, _Tp>(x, parameters.scale, 0.5);
          }

          if (y_index < maxVoxelDataSize) {
            voxelData[y_index] = quantify<float, _Tp>(y, parameters.scale, 0.5);
          }

          if (z_index < maxVoxelDataSize) {
            voxelData[z_index] = quantify<float, _Tp>(z, parameters.scale, 0.5);
          }

          if (i_index < maxVoxelDataSize) {
            voxelData[i_index] =
                quantify<float, _Tp>(intensity, parameters.scale, 0.5);
          }

          ++voxel_points_num_lut_ptr_->at(voxelPointIndex);
        }
      }
    }};

    uint32_t threads{parameters.multi_threads};
    if (threads > 1) {
      std::vector<std::shared_future<void>> futureList;
      for (size_t i = 0; i < threads; i++) {
        std::shared_future<void> f{std::async(
            std::launch::async, doVoxelFunc, i * pointCloudNum / threads,
            (i + 1) * pointCloudNum / threads)};
        futureList.emplace_back(f);
      }

      for (auto& future : futureList) {
        future.wait();
      }
    } else {
      doVoxelFunc(0, pointCloudNum);
    }

    return true;
  }

  template <std::uint32_t maxPointSize>
  bool Sparsity(PointCloud<PointXYZIT, maxPointSize> const& pclIn,
                VoxelParameters const& parameters,
                std::vector<PointCloudSparseData<float>>& voxelData,
                std::string& err) noexcept {
    auto const& voxelSize{parameters.voxel_size};
    auto const& netRange{parameters.net_range};
    auto const& m1{parameters.mp.at(0)};
    auto const& m2{parameters.mp.at(1)};
    auto const& maxPointsPerVoxel{parameters.max_points_per_voxel};

    auto const& pointCloudNum{pclIn.GetPointNum()};
    // 初始化
    auto const& ret{init(parameters, pointCloudNum, err)};
    if (!ret) {
      return false;
    }

    // voxel
    float vx{voxelSize[0] / 2 + netRange[0]};
    float vy{voxelSize[1] / 2 + netRange[1]};
    float vz{voxelSize[2] / 2 + netRange[2]};
    float z_coe{voxelSize[2] / 2};

    std::vector<int32_t> grid(voxelSize.size(), 0);

    // feature map h
    auto const& h{static_cast<int32_t>(grid_size_[1] * m2)};
    // feature map w
    auto const& w{static_cast<int32_t>(grid_size_[0] * m1)};
    // feature map stride
    auto const& stride{h * w};

    auto const& pointCloud{pclIn};
    auto doVoxelFunc{[&](int32_t startIndex, int32_t stopIndex) {
      for (auto j{startIndex}; j < stopIndex; ++j) {
        auto const& selectIndex{random_array_ptr_->at(j)};
        if (selectIndex < pointCloudNum) {
          auto& point{pointCloud.GetPoint(selectIndex)};
          grid[0] = floor((point.x - netRange[0]) / voxelSize[0]);
          if (grid[0] < 0 || grid[0] >= grid_size_[0]) {
            continue;
          }
          grid[1] = floor((point.y - netRange[1]) / voxelSize[1]);
          if (grid[1] < 0 || grid[1] >= grid_size_[1]) {
            continue;
          }
          grid[2] = floor((point.z - netRange[2]) / voxelSize[2]);
          if (grid[2] < 0 || grid[2] >= grid_size_[2]) {
            continue;
          }

          auto const& voxelPointIndex{
              gen_index(grid[0], grid[1], grid_size_[0], grid_size_[1])};

          if (voxelPointIndex >= voxel_points_num_lut_ptr_->size()) {
            continue;
          }

          auto const& voxelPointNumber{
              voxel_points_num_lut_ptr_->at(voxelPointIndex).load()};

          if (voxelPointNumber >= maxPointsPerVoxel) {
            continue;
          }

          // position in single feature map
          std::int32_t pos{(grid[1] * m1 + voxelPointNumber % m1) * h +
                           (grid[0] * m2 + voxelPointNumber / m2)};

          // feature value
          auto const& x{(point.x - (grid[0] * voxelSize[0] + vx)) /
                        voxelSize[0] * 2.0};
          auto const& y{(point.y - (grid[1] * voxelSize[1] + vy)) /
                        voxelSize[1] * 2.0};
          auto const& z{(point.z - vz) / z_coe};
          auto const& intensity{(static_cast<float>(point.intensity) - 127.5) /
                                127.5};

          auto const& x_b{point.x >= 0 ? 1 : -1};
          auto const& y_b{point.y >= 0 ? 1 : -1};

          PointCloudSparseData<float> sparseData;
          sparseData.position = pos;
          sparseData.feature.x = x;
          sparseData.feature.y = y;
          sparseData.feature.z = z;
          sparseData.feature.intensity = intensity;
          sparseData.feature.x_b = x_b;
          sparseData.feature.y_b = y_b;
          voxelData.emplace_back(sparseData);

          ++voxel_points_num_lut_ptr_->at(voxelPointIndex);
        }
      }
    }};

    uint32_t threads{parameters.multi_threads};
    if (threads > 1) {
      std::vector<std::shared_future<void>> futureList;
      for (size_t i = 0; i < threads; i++) {
        std::shared_future<void> f{std::async(
            std::launch::async, doVoxelFunc, i * pointCloudNum / threads,
            (i + 1) * pointCloudNum / threads)};
        futureList.emplace_back(f);
      }

      for (auto& future : futureList) {
        future.get();
      }
    } else {
      doVoxelFunc(0, pointCloudNum);
    }

    return true;
  }

  /// @brief
  void Reset() noexcept;

 protected:
  /// @brief
  /// @param para
  /// @param pointcloud_number
  /// @param err
  /// @return
  bool init(VoxelParameters const& para, uint32_t pointcloud_number,
            std::string& err) noexcept;

  /// @brief
  /// @param coor_x
  /// @param coor_y
  /// @param h
  /// @param w
  /// @return
  uint32_t gen_index(uint32_t coor_x, uint32_t coor_y, uint32_t grid_size_x,
                     uint32_t grid_size_y) noexcept {
    return coor_x + grid_size_x * coor_y;
  }

 private:
  /// @brief random array
  std::shared_ptr<std::vector<uint32_t>> random_array_ptr_;
  /// @brief  init call once flag
  std::once_flag init_flag_;
  /// @brief number of per voxel lookup talble
  std::shared_ptr<std::vector<std::atomic_uint8_t>> voxel_points_num_lut_ptr_;
  /// @brief grid size
  std::vector<uint32_t> grid_size_;
};
}  // namespace perception
}  // namespace autodrive

namespace boost::json {
inline void tag_invoke(value_from_tag, value& jv,
                       autodrive::perception::Voxel::VoxelParameters const& v) {
  jv = {{"net_range", value_from(v.net_range)},
        {"voxel_size", value_from(v.voxel_size)},
        {"mp", value_from(v.mp)},
        {"sd", v.sd},
        {"max_points_per_voxel", v.max_points_per_voxel},
        {"scale", v.scale},
        {"multi_threads", v.multi_threads},
        {"is_shuffle", v.is_shuffle}};
}

inline autodrive::perception::Voxel::VoxelParameters tag_invoke(
    value_to_tag<autodrive::perception::Voxel::VoxelParameters>,
    value const& jv) {
  auto const& obj = jv.as_object();
  return autodrive::perception::Voxel::VoxelParameters{
      boost::json::value_to<std::vector<float>>(obj.at("net_range")),
      boost::json::value_to<std::vector<float>>(obj.at("voxel_size")),
      boost::json::value_to<std::vector<int32_t>>(obj.at("mp")),
      boost::json::value_to<int32_t>(obj.at("sd")),
      boost::json::value_to<int32_t>(obj.at("max_points_per_voxel")),
      boost::json::value_to<int32_t>(obj.at("scale")),
      boost::json::value_to<uint16_t>(obj.at("multi_threads")),
      boost::json::value_to<bool>(obj.at("is_shuffle"))};
}
}  // namespace boost::json

#endif  // !VOXEL_H_
